package com.example;

import redis.clients.jedis.Connection;
import redis.clients.jedis.Jedis;
import redis.clients.jedis.commands.ProtocolCommand;
import redis.clients.jedis.util.SafeEncoder;

import java.util.List;

// Redis 4.0 提供了一个限流Redis 模块， 它日叫Redis-Cell 。该模块也使用了漏斗算法，
// 并提供了原子的限流指令CL.THROTTLE
public class RedisCellRateLimiter {

    private Jedis jedis;

    public RedisCellRateLimiter(Jedis jedis) {
        this.jedis = jedis;
    }

    /**
     * 判断用户的行为是否应该被允许
     * @param userId 用户id
     * @param actionId 行为
     * @param maxBurst 漏斗的最大容量
     * @param countPerPeriod 每个period漏水量 （countPerPeriod/period=漏水的速率）
     * @param period 时间段，单位为秒
     * @return 允许则返回true，否则返回false.
     */
    public boolean isActionAllowed(String userId, String actionId,
                                   Integer maxBurst,
                                   Integer countPerPeriod, Integer period){

        String key = String.format("%s:%s", userId, actionId);

        Connection client = jedis.getClient();

        //命令格式： CL.THROTTLE <key> <max_burst> <count per period> <period> [<quantity>]
        client.sendCommand(CellCommand.CLTHROTTLE, key, maxBurst.toString(), countPerPeriod.toString(), period.toString());
        // 返回值格式：
        //1) (integer) 0   # 0 表示允许， 1 表示拒绝
        //2) (integer) 16  # 漏斗容量capacity  (max_burst + 1)
        //3) (integer) 15  # 漏斗剩余空间left_quota
        //4) (integer) -1  # 如采被拒绝了，需要多长时间后再试（漏斗有空间了，单位秒）
        //5) (integer) 2   # 多长时间后，漏斗完全空出来（ left_quota==capacity ，单位秒）
        List<Long> replay = client.getIntegerMultiBulkReply();

        return replay.get(0) == 0;
    }

    // 因为CL.THROTTLE是扩展命令，所以要先继承ProtocolCommand扩展命令
    public enum CellCommand implements ProtocolCommand{
        CLTHROTTLE("CL.THROTTLE");

        private byte[] raw;

        private CellCommand(String cmd){
            this.raw = SafeEncoder.encode(cmd);
        }

        @Override
        public byte[] getRaw() {
            return raw;
        }
    }
}
